# -*- coding: utf-8 -*-
"""
Created on Tue Oct 24 09:44:25 2017
修改xml中错误的节点,检查object个数
@author: aia_xml
"""

import xml.etree.ElementTree as ET
import os
import shutil

'''计算文件夹下文件个数'''


def fileCountIn(dir):
    return sum([len(files) for root, dirs, files in os.walk(dir)])


def main():
    xml_dir = '/home/dataset/zhangyan/15_yolov3_4/Annotations'
    save_dir = '/home/ccj/dataset/test-sets/pedestrian'

    xml_Count = fileCountIn(xml_dir)
    print('xml Count in dir:', xml_Count)
    chg_num = 0
    copied_img = 0
    object_dicts = {}
    for root, dirs, files in os.walk(xml_dir):
        for file in files:
            xml_file = os.path.join(xml_dir, file)
            # change_xml_node(xml_file)
            # objects = find_objects(xml_file)

            tree = ET.parse(xml_file)
            root = tree.getroot()
            for obj in root.iter('object'):
                if obj.find('name').text in object_dicts:
                    object_dicts[obj.find('name').text] += 1
                else:
                    object_dicts.update({obj.find('name').text: 1})


            chg_num = chg_num + 1
            if chg_num % 1000 == 0:
                print('completed ', chg_num * 100 / xml_Count, '% ...')
    print('completed ', chg_num * 100 / xml_Count, '% ...')
    print('total:', chg_num)
    print('object_dicts: ', object_dicts)


if __name__ == '__main__':
    main()
